/* Copyright 2018 The TensorFlow Authors. All Rights Reserved.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/

#include <stdint.h>

#include <algorithm>
#include <string>
#include <vector>

#include "tensorflow/lite/c/builtin_op_data.h"
#include "tensorflow/lite/c/common.h"
#include "tensorflow/lite/kernels/internal/tensor_ctypes.h"
#include "tensorflow/lite/kernels/kernel_util.h"
#include "tensorflow/lite/micro/kernels/kernel_util.h"

namespace tflite {
namespace ops {
namespace micro {
namespace slice {

enum KernelType {
  kReference,
  kGenericOptimized,
};

constexpr int kInputTensor = 0;
constexpr int kBeginTensor = 1;
constexpr int kSizeTensor = 2;
constexpr int kOutputTensor = 0;

// This Op only supports 1-5D cases and since we use the optimized ops 5D
// implementation, the 1-4D tensors are mapped to 5D.
const int kMaxDim = 5;

#define DIMS_SIZE(x) (x->dims->size)
#define DIMS_DATA(x) (x->dims->data)

template <typename T>
TfLiteStatus Slice(TfLiteContext* context, TfLiteNode* node,
                   const TfLiteEvalTensor* input, SliceParams* params){

	const T* input_ptr = tflite::micro::GetTensorData<T>(input);
	TfLiteEvalTensor* t = tflite::micro::GetEvalOutput(context, node, kOutputTensor);
	T* output_data = tflite::micro::GetTensorData<T>(t);
					   
	uint32_t strides[5] = { 1, 0, 0, 1, 1 };
    uint32_t strides_mul = 1;
    // start from the second one, the first one is 1
    int32_t* data = DIMS_DATA(input) + DIMS_SIZE(input) - 1;
    for (int idx = 0; data >= DIMS_DATA(input) + 1; data--, idx++) {
        strides_mul *= (*data);
        strides[kMaxDim - idx - 2] = strides_mul;
    }

    uint32_t new_nn, new_n, new_h, new_w, new_c;
    uint32_t offset_nn, offset_n, offset_h, offset_w, offset_c;
	new_nn = new_n = new_h = new_w = new_c = 1;
	switch(DIMS_SIZE(t)){
		case 5:
			new_nn = DIMS_DATA(t)[DIMS_SIZE(t) - 5];
			// no break;
		case 4:
			new_n = DIMS_DATA(t)[DIMS_SIZE(t) - 4];
			// no break;
		case 3:
			new_h = DIMS_DATA(t)[DIMS_SIZE(t) - 3];
			// no break;
		case 2:
			new_w = DIMS_DATA(t)[DIMS_SIZE(t) - 2];
			new_c = DIMS_DATA(t)[DIMS_SIZE(t) - 1];
			break;
		default:
			break;
	}

    offset_nn = params->begin[0];
    offset_n = params->begin[1];
    offset_h = params->begin[2];
    offset_w = params->begin[3];
    offset_c = params->begin[4];

    uint32_t array_tick = 0;
    for (uint32_t i = 0;i < new_nn;i++) {
        uint32_t strides_nn = strides[0] * (i + offset_nn);
        for (uint32_t j = 0;j < new_n;j++) {
            uint32_t strides_n = strides[1] * (j + offset_n);
            for (uint32_t k = 0;k < new_h;k++) {
                uint32_t strides_h = strides[2] * (k + offset_h);
                for (uint32_t l = 0;l < new_w;l++) {
                    uint32_t strides_w = strides[3] * (l + offset_w);
                    for (uint32_t m = 0;m < new_c;m++) {
                        uint32_t strides_c = strides[4] * (m + offset_c);
                        output_data[array_tick++] = input_ptr[strides_c + strides_w + strides_h + strides_n + strides_nn];
                    }
                }
            }
        }
    }

	return kTfLiteOk;
				   
}
template <typename T>
TfLiteStatus CalculateOutputShapeVector(TfLiteContext* context,
                                        const TfLiteTensor* input,
                                        const TfLiteTensor* begin,
                                        const TfLiteTensor* size,
                                        std::vector<int>* output_shape_vector) {
  for (int idx = 0; idx < NumDimensions(input); ++idx) {
    T size_value = GetTensorData<T>(size)[idx];
    if (size_value < 0) {
      if (size_value != -1) {
        context->ReportError(context, "Invalid size.");
        return kTfLiteError;
      }
      size_value = SizeOfDimension(input, idx) - GetTensorData<T>(begin)[idx];
    } else {
      if (SizeOfDimension(input, idx) <
          GetTensorData<T>(begin)[idx] + size_value) {
        context->ReportError(context, "Invalid begin and size.");
        return kTfLiteError;
      }
    }
    output_shape_vector->push_back(static_cast<int>(size_value));
  }
  return kTfLiteOk;
}

template <typename T>
void GetBeginAndSizeVectors(int dimensions, const TfLiteEvalTensor* begin,
                            const TfLiteEvalTensor* size, std::vector<int>* begins,
                            std::vector<int>* sizes) {
  for (int idx = 0; idx < dimensions; ++idx) {
	begins->push_back(tflite::micro::GetTensorData<T>(begin)[idx]);
    sizes->push_back(tflite::micro::GetTensorData<T>(size)[idx]);
  }
}

TfLiteStatus ResizeOutputShape(TfLiteContext* context,
                               const TfLiteTensor* input,
                               const TfLiteTensor* begin,
                               const TfLiteTensor* size, TfLiteTensor* output) {
  std::vector<int> output_shape_vector;

  if (begin->type == kTfLiteInt32) {
    TF_LITE_ENSURE_STATUS(CalculateOutputShapeVector<int32_t>(
        context, input, begin, size, &output_shape_vector));
  } else if (begin->type == kTfLiteInt64) {
    TF_LITE_ENSURE_STATUS(CalculateOutputShapeVector<int64_t>(
        context, input, begin, size, &output_shape_vector));
  } else {
    context->ReportError(
        context, "Type %d is currently not supported by Slice.", begin->type);
    return kTfLiteError;
  }

  TfLiteIntArray* output_shape = output->dims;
  std::copy(output_shape_vector.begin(), output_shape_vector.end(),
            output_shape->data);
  
  return kTfLiteOk;
}

TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
  TF_LITE_ENSURE_EQ(context, NumInputs(node), 3);
  TF_LITE_ENSURE_EQ(context, NumOutputs(node), 1);

  const TfLiteTensor* input;
  TF_LITE_ENSURE_OK(context, GetInputSafe(context, node, kInputTensor, &input));
  const TfLiteTensor* begin;
  TF_LITE_ENSURE_OK(context, GetInputSafe(context, node, kBeginTensor, &begin));
  const TfLiteTensor* size;
  TF_LITE_ENSURE_OK(context, GetInputSafe(context, node, kSizeTensor, &size));
  TfLiteTensor* output;
  TF_LITE_ENSURE_OK(context,
                    GetOutputSafe(context, node, kOutputTensor, &output));

  // Ensure validity of input tensor and its dimension.
  TF_LITE_ENSURE_TYPES_EQ(context, input->type, output->type);
  TF_LITE_ENSURE(context,
                 begin->type == kTfLiteInt32 || begin->type == kTfLiteInt64);
  TF_LITE_ENSURE(context,
                 size->type == kTfLiteInt32 || size->type == kTfLiteInt64);
  TF_LITE_ENSURE_EQ(context, NumDimensions(begin), 1);
  TF_LITE_ENSURE_EQ(context, NumDimensions(size), 1);
  TF_LITE_ENSURE_EQ(context, NumElements(begin), NumElements(size));
  TF_LITE_ENSURE_MSG(context, NumDimensions(input) <= kMaxDim,
                     "Slice op only supports 1D-5D input arrays.");

  // Postpone allocation of output if any of the indexing tensors is not
  // constant
  if (!(IsConstantTensor(begin) && IsConstantTensor(size))) {
    SetTensorToDynamic(output);
    return kTfLiteOk;
  }

  return ResizeOutputShape(context, input, begin, size, output);
}

TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) {
  const TfLiteEvalTensor* input = 
	  tflite::micro::GetEvalInput(context, node, kInputTensor);
  const TfLiteEvalTensor* begin = 
 	  tflite::micro::GetEvalInput(context, node, kBeginTensor);
  const TfLiteEvalTensor* size =
 	  tflite::micro::GetEvalInput(context, node, kSizeTensor);
  const TfLiteEvalTensor* output =
 	  tflite::micro::GetEvalOutput(context, node, kOutputTensor);

  if (IsDynamicTensor(GetOutput(context, node, kOutputTensor))) {
    TF_LITE_ENSURE_OK(context,
                      ResizeOutputShape(context, 
									    (const TfLiteTensor*)GetInput(context, node, kInputTensor), 
	                                    (const TfLiteTensor*)GetInput(context, node, kBeginTensor), 
	                                    (const TfLiteTensor*)GetInput(context, node, kSizeTensor), 
	                                     GetOutput(context, node, kOutputTensor)));
  }
  std::vector<int> begins;
  begins.reserve(kMaxDim);
  std::vector<int> sizes;
  sizes.reserve(kMaxDim);

  for (int i = input->dims->size; i < kMaxDim; ++i) {
    begins.push_back(0);
    sizes.push_back(1);
  }
  
  // The Slice op implementation only accepts 5-D sizes. That constraint is, for
  // the present, maintained here.
  //
  // The dimensions in the kernel used to be in reverse-order, and TFLite
  // arranged the begins and sizes vectors accordingly. This macro incorporates
  // the needed reversing.
  if (begin->type == kTfLiteInt32) {
    GetBeginAndSizeVectors<int32_t>(input->dims->size, begin, size, &begins,
                                    &sizes);
  } else if (begin->type == kTfLiteInt64) {
    GetBeginAndSizeVectors<int64_t>(input->dims->size, begin, size, &begins,
                                    &sizes);
  } else {
    context->ReportError(
        context, "Type %d is currently not supported by Slice.", begin->type);
    return kTfLiteError;
  }
  
  TF_LITE_ENSURE_EQ(context, begins.size(), kMaxDim);                        
  TF_LITE_ENSURE_EQ(context, sizes.size(), kMaxDim);                         
  tflite::SliceParams/*use tflite type instead of tflm's*/ op_params;                                             
  op_params.begin_count = kMaxDim;                                           
  op_params.size_count = kMaxDim;                                            
  for (int i = 0; i < kMaxDim; ++i) {                                        
    op_params.begin[i] = begins[i];                                          
    op_params.size[i] = sizes[i];                                            
  }  

  #define TF_LITE_SLICE(type)   Slice<type>(context, node, input, &op_params);
  
  switch (input->type) {
    case kTfLiteFloat32:
      TF_LITE_SLICE(float);
      break;
    case kTfLiteInt32:
      TF_LITE_SLICE(int32_t);
      break;
    case kTfLiteInt64:
      TF_LITE_SLICE(int64_t);
      break;
    case kTfLiteInt8:
      TF_LITE_SLICE(int8_t);
      break;
    case kTfLiteInt16:
      TF_LITE_SLICE(int16_t);
      break;
    case kTfLiteUInt8:
      TF_LITE_SLICE(uint8_t);
      break;
    case kTfLiteBool:
      TF_LITE_SLICE(bool);
      break;
    case kTfLiteString:
	  TF_LITE_SLICE(std::string);
      break;
    default:
      context->ReportError(
          context, "Type %d is currently not supported by Slice.", input->type);
      return kTfLiteError;
  }
  #undef TF_LITE_SLICE
  return kTfLiteOk;
}
}  // namespace slice

TfLiteRegistration Register_SLICE() {
  return {/*init=*/nullptr,
          /*free=*/nullptr,
		  /*prepare=*/slice::Prepare,
          /*invoke=*/slice::Eval,
          /*profiling_string=*/nullptr,
          /*builtin_code=*/0,
          /*custom_name=*/nullptr,
          /*version=*/0};
}

}  // namespace micro
}  // namespace ops
}  // namespace tflite
